import os

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import Dataset, DataLoader


class CustomerDataLoader(Dataset):
    def __init__(self, img_dir, mask_dir, transform=None):
        """
        初始化数据集
        :param img_dir: 存放图片的目录
        :param mask_dir: 存放mask的目录
        :param transform: 预处理操作（如转换为Tensor等）
        """
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform

        # 获取所有图片文件名（假设图片和mask文件的名称相同，扩展名不同）
        self.img_files = [f for f in os.listdir(img_dir) if f.endswith('.bmp')]

    def __len__(self):
        """
        返回数据集大小
        """
        return len(self.img_files)

    def __getitem__(self, idx):
        """
        返回图像和对应的mask
        :param idx: 数据索引
        """
        # 获取图像文件名
        img_name = self.img_files[idx]

        # 加载图像
        img_path = os.path.join(self.img_dir, img_name)
        img = Image.open(img_path).convert('RGB')

        # 获取mask文件名（假设mask文件名与图像相同，只是保存在mask目录中）
        mask_name = img_name.replace('.bmp', '_anno.bmp')
        mask_path = os.path.join(self.mask_dir, mask_name)
        mask = Image.open(mask_path).convert('L')  # 假设mask是灰度图

        # 将mask中非零值设置为1
        mask = np.array(mask)  # 转换为numpy数组
        mask[mask > 0] = 1  # 非零值设置为1

        # 将mask转换回PIL.Image后再进行transform
        mask = Image.fromarray(mask)  # 转换回PIL.Image

        # 应用必要的转换
        if self.transform:
            img = self.transform(img)

            # 对mask应用transform时，需要确保mask的值在[0, 1]范围内
            # 使用 lambda 函数来确保mask的值是[0, 1]之间
            mask = self.transform(mask)
            mask = (mask != 0).float()  # 将mask值转换为0和1
            # print(torch.max(mask), torch.min(mask))
        return img, mask

    def display_image(self, idx):
        """
        展示图像和mask
        :param idx: 数据索引
        """
        img, mask = self[idx]

        # 使用matplotlib显示图像和mask
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
        ax1.imshow(np.array(img))
        ax1.set_title("Image")

        ax2.imshow(np.array(mask), cmap='gray')
        ax2.set_title("Mask")
        plt.show()


# 使用示例
if __name__ == "__main__":
    img_dir = "/Volumes/For_Mac/Download/GlasDateset/train/images"  # 图片文件夹路径
    mask_dir = "/Volumes/For_Mac/Download/GlasDateset/train/masks"  # mask文件夹路径

    # 自定义的transform，包含resize和转换为tensor操作
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 调整图像大小为224x224
        transforms.ToTensor(),  # 转换为Tensor格式
    ])
    dataset = CustomerDataLoader(img_dir, mask_dir, transform=transform)

    # 查看数据集中的第一张图像和mask
    # for i in range(1):
    #     dataset.display_image(i)

    # 使用DataLoader加载数据
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)
    for img, mask in dataloader:
        # 你可以在这里进行批处理
        # print(img.shape, mask.shape)
        pass
